# -*- coding: utf-8 -*-
"""
@File: validate_mode.py
@Time: 2021/12/18 18:57
@Author: 鹄望潇湘
@desc:
使用固定的前30帧的所有ID的数据进行训练30 epoch，使用ID为75-123的数据进行测试，rank-1结果如下：
====平均值====
NM: 89.161 BG: 86.823 CL: 81.354
====详细值====
CL: 84.69 75.51 80.61 69.39 89.80 84.69 84.69 78.57 78.57 79.59 88.78
BG: 93.88 84.54 82.47 69.79 85.71 92.86 92.86 80.61 87.76 88.66 95.92
NM: 95.55 87.50 84.32 83.22 92.86 93.20 91.84 82.29 86.21 88.97 94.83

使用随机的30帧，所有ID的数据进行训练100 epoch，使用ID为75-123的数据进行测试，rank-1结果如下:
====平均值====
NM: 98.775 BG: 98.321 CL: 97.403
====详细值====
CL: 96.94 97.96 98.98 98.98 96.94 97.96 98.98 97.96 98.98 92.86 94.90
BG: 95.92 97.94 94.85 98.96 100.00 97.96 98.98 98.98 100.00 98.97 98.98
NM: 99.66 99.65 98.26 97.20 99.32 98.98 99.32 98.96 99.66 97.93 97.59


"""

import torch
from torch.utils.data.dataloader import DataLoader
from tqdm import tqdm

from dataset.GaitMChannelDataSet import GaitMChannelDataLoader, SingleGaitSequence, GaitDataSet
import numpy as np

# test_id = ["075", "076", "077", "078", "079", "080", "081", "082", "083", "084", "085", "086",
#           "087", "088", "089", "090", "091", "092", "093", "094", "095", "096", "097", "098"]
test_id = [index for index in range(75, 124)]
angle_type = {'000': 0, '018': 1, '036':2, '054': 3, '072': 4, '090': 5,
              '108': 6, '126': 7, '144': 8, '162': 9,'180': 10}
types_type = {'cl-01': 0, 'cl-02': 0,
              'bg-01': 1, 'bg-02': 1,
              'nm-01': 2, 'nm-02': 2, 'nm-03': 2, 'nm-04': 2, 'nm-05': 2, 'nm-06': 2}
correct_metric = np.zeros(shape=(3, 11), dtype=np.int)
sum_metric = np.zeros(shape=(3, 11), dtype=np.int)

data_loader = GaitMChannelDataLoader("E:\\DataSet\\CASIA-B-preprocessed", True)
datas = data_loader.get_partial_data_with_id(test_id)
dataset = GaitDataSet(datas)
dataloader = DataLoader(dataset, batch_size=1, num_workers=0, shuffle=False)

model = torch.load("./checkpoint_29.pth")
gpu = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

top_k = 1

for id, item in tqdm(enumerate(dataloader)):
    label = item[1]
    angle = item[2][0]
    types = item[3][0]
    output = model(item[0].float().to(gpu))
    _, idx = torch.sort(output, dim=1, descending=True)
    idx = idx[0, 0: top_k].cpu()

    sum_metric[types_type[types], angle_type[angle]] += 1
    if label in idx:
        correct_metric[types_type[types], angle_type[angle]] += 1
    pass

percent_metric = correct_metric*100/sum_metric
print("====平均值====")
average = np.mean(percent_metric,axis=1)
print("NM: {:.3f} BG: {:.3f} CL: {:.3f}".format(average[2], average[1], average[0]))
print("====详细值====")
for row in range(3):
   strs = ""
   if row==0:
       strs= strs + "CL: "
   if row == 1:
       strs = strs + "BG: "
   if row ==2:
       strs = strs + "NM: "
   for col in range(11):
       strs = strs + "{:.2f} ".format(percent_metric[row, col])
   print(strs)



